# sources/kafka.py
from datetime import datetime
import json
from confluent_kafka import Consumer, KafkaException


# Source算子
class KafkaSource:
    def __init__(self, bootstrap_servers, topic, group_id='wordcount-group'):
        self.conf = {
            'bootstrap.servers': bootstrap_servers,
            'group.id': group_id,
            'auto.offset.reset': 'earliest',
            'enable.auto.commit': False  # 禁用自动提交
        }
        self.topic = topic
        self.consumer = Consumer(self.conf)
        self.consumer.subscribe([self.topic])


    def __iter__(self):
        """持续消费Kafka数据"""
        try:
            while True:
                msg = self.consumer.poll(1.0)
                if msg is None:
                    continue
                if msg.error():
                    raise KafkaException(msg.error())
                
                data = json.loads(msg.value())
                data['timestamp'] = datetime.fromisoformat(data['timestamp'])
                yield data
        except KeyboardInterrupt:
            print("\n消费已停止")
        finally:
            self.consumer.close()
from confluent_kafka import Consumer, KafkaError

class KafkaSource2:
    def __init__(self, bootstrap_servers, topic, group_id='wordcount-group'):
        self.conf = {
            'bootstrap.servers': bootstrap_servers,
            'group.id': group_id,
            'auto.offset.reset': 'earliest',
            'enable.auto.commit': False
        }
        self.topic = topic
        self.consumer = Consumer(self.conf)
        self.consumer.subscribe([self.topic])

    def __iter__(self):
        while True:
            msg = self.consumer.poll(1.0)
            if msg is None:
                continue
            if msg.error():
                if msg.error().code() == KafkaError._PARTITION_EOF:
                    continue
                else:
                    print(msg.error())
                    break
            yield json.loads(msg.value().decode('utf-8'))
            # 手动提交偏移量
            self.consumer.commit(message=msg)